#
#  tsne_torch.py
#
# Implementation of t-SNE in pytorch. The implementation was tested on pytorch
# > 1.0, and it requires Numpy to read files. In order to plot the results,
# a working installation of matplotlib is required.
#
#
# The example can be run by executing: `python tsne_torch.py`
#
#
#  Created by Xiao Li on 23-03-2020.
#  Copyright (c) 2020. All rights reserved.
import os

import argparse
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import shutil
import torch
from tqdm import tqdm

base = 'shrec14'

if base == 'part-shrec14':
    classes = 48
    sk_load = os.path.join('.', 'labels', 'PART-SHREC14-EDGE', 'sk_edge.hdf5')
    cad_load = os.path.join('.', 'labels', 'PART-SHREC14-EDGE', 'cad_edge.hdf5')
    res_load = os.path.join('.', 'pth', 'part-shrec14', 'retrieval_result.pth')
    save_dir = os.path.join('.', 'retrieval_result', 'part-shrec14')
elif base == 'shrec14':
    classes = 171
    sk_load = os.path.join('.', 'labels', 'SHREC14-EDGE', 'sk_edge.hdf5')
    cad_load = os.path.join('.', 'labels', 'SHREC14-EDGE', 'cad_edge.hdf5')
    res_load = os.path.join('.', 'pth', 'shrec14', 'retrieval_result.pth')
    save_dir = os.path.join('.', 'retrieval_result', 'shrec14')
elif base == 'shrec13':
    classes = 90
    sk_load = os.path.join('.', 'labels', 'SHREC13-EDGE', 'sk_edge.hdf5')
    cad_load = os.path.join('.', 'labels', 'SHREC13-EDGE', 'cad_edge.hdf5')
    res_load = os.path.join('.', 'pth', 'shrec13', 'retrieval_result.pth')
    save_dir = os.path.join('.', 'retrieval_result', 'shrec13')

parser = argparse.ArgumentParser()
parser.add_argument("--cuda", type=int, default=1, help="if use cuda accelarate")

opt = parser.parse_args()

if opt.cuda:
    print("set use cuda")
    torch.set_default_tensor_type(torch.cuda.DoubleTensor)
else:
    torch.set_default_tensor_type(torch.DoubleTensor)


def Hbeta_torch(D, beta=1.0):
    P = torch.exp(-D.clone() * beta)

    sumP = torch.sum(P)

    H = torch.log(sumP) + beta * torch.sum(D * P) / sumP
    P = P / sumP

    return H, P


def x2p_torch(X, tol=1e-5, perplexity=30.0):
    """
        Performs a binary search to get P-values in such a way that each
        conditional Gaussian has the same perplexity.
    """

    # Initialize some variables
    print("Computing pairwise distances...")
    (n, d) = X.shape

    sum_X = torch.sum(X * X, 1)
    D = torch.add(torch.add(-2 * torch.mm(X, X.t()), sum_X).t(), sum_X)

    P = torch.zeros(n, n)
    beta = torch.ones(n, 1)
    logU = torch.log(torch.tensor([perplexity]))
    n_list = [i for i in range(n)]

    # Loop over all datapoints
    for i in range(n):

        # Print progress
        if i % 500 == 0:
            print("Computing P-values for point %d of %d..." % (i, n))

        # Compute the Gaussian kernel and entropy for the current precision
        # there may be something wrong with this setting None
        betamin = None
        betamax = None
        Di = D[i, n_list[0:i] + n_list[i + 1:n]]

        (H, thisP) = Hbeta_torch(Di, beta[i])

        # Evaluate whether the perplexity is within tolerance
        Hdiff = H - logU
        tries = 0
        while torch.abs(Hdiff) > tol and tries < 50:

            # If not, increase or decrease precision
            if Hdiff > 0:
                betamin = beta[i].clone()
                if betamax is None:
                    beta[i] = beta[i] * 2.
                else:
                    beta[i] = (beta[i] + betamax) / 2.
            else:
                betamax = beta[i].clone()
                if betamin is None:
                    beta[i] = beta[i] / 2.
                else:
                    beta[i] = (beta[i] + betamin) / 2.

            # Recompute the values
            (H, thisP) = Hbeta_torch(Di, beta[i])

            Hdiff = H - logU
            tries += 1

        # Set the final row of P
        P[i, n_list[0:i] + n_list[i + 1:n]] = thisP

    # Return final P-matrix
    return P


def pca_torch(X, no_dims=50):
    print("Preprocessing the data using PCA...")
    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    (l, M) = torch.eig(torch.mm(X.t(), X), True)
    # split M real
    for i in range(d):
        if l[i, 1] != 0:
            M[:, i + 1] = M[:, i]
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims])
    return Y


def tsne(X, no_dims=2, initial_dims=50, perplexity=30.0):
    """
        Runs t-SNE on the dataset in the NxD array X to reduce its
        dimensionality to no_dims dimensions. The syntaxis of the function is
        `Y = tsne.tsne(X, no_dims, perplexity), where X is an NxD NumPy array.
    """

    # Check inputs
    if isinstance(no_dims, float):
        print("Error: array X should not have type float.")
        return -1
    if round(no_dims) != no_dims:
        print("Error: number of dimensions should be an integer.")
        return -1

    # Initialize variables
    X = pca_torch(X, initial_dims)
    (n, d) = X.shape
    max_iter = 1000
    initial_momentum = 0.5
    final_momentum = 0.8
    eta = 500
    min_gain = 0.01
    Y = torch.randn(n, no_dims)
    dY = torch.zeros(n, no_dims)
    iY = torch.zeros(n, no_dims)
    gains = torch.ones(n, no_dims)

    # Compute P-values
    P = x2p_torch(X, 1e-5, perplexity)
    P = P + P.t()
    P = P / torch.sum(P)
    P = P * 4.  # early exaggeration
    print("get P shape", P.shape)
    P = torch.max(P, torch.tensor([1e-21]))

    # Run iterations
    for iter in range(max_iter):

        # Compute pairwise affinities
        sum_Y = torch.sum(Y * Y, 1)
        num = -2. * torch.mm(Y, Y.t())
        num = 1. / (1. + torch.add(torch.add(num, sum_Y).t(), sum_Y))
        num[range(n), range(n)] = 0.
        Q = num / torch.sum(num)
        Q = torch.max(Q, torch.tensor([1e-12]))

        # Compute gradient
        PQ = P - Q
        for i in range(n):
            dY[i, :] = torch.sum((PQ[:, i] * num[:, i]).repeat(no_dims, 1).t() * (Y[i, :] - Y), 0)

        # Perform the update
        if iter < 20:
            momentum = initial_momentum
        else:
            momentum = final_momentum

        gains = (gains + 0.2) * ((dY > 0.) != (iY > 0.)).double() + (gains * 0.8) * ((dY > 0.) == (iY > 0.)).double()
        gains[gains < min_gain] = min_gain
        iY = momentum * iY - eta * (gains * dY)
        Y = Y + iY
        Y = Y - torch.mean(Y, 0)

        # Compute current value of cost function
        if (iter + 1) % 10 == 0:
            C = torch.sum(P * torch.log(P / Q))
            print("Iteration %d: error is %f" % (iter + 1, C))

        # Stop lying about P-values
        if iter == 100:
            P = P / 4.

    # Return solution
    return Y


def save_tsne_joint_result(sk_f, sk_l, cad_f, cad_l, filename):
    assert (len(sk_f[:, 0]) == len(sk_f[:, 1]))
    assert (len(sk_f) == len(sk_l))
    assert (len(cad_f[:, 0]) == len(cad_f[:, 1]))
    assert (len(cad_f) == len(cad_l))

    with torch.no_grad():
        sk_f = tsne(sk_f, 2, 50, 20.0)
        cad_f = tsne(cad_f, 2, 50, 20.0)

    if opt.cuda:
        sk_f = sk_f.cpu().numpy()
        cad_f = cad_f.cpu().numpy()
    sk_l = sk_l.cpu().numpy()
    cad_l = cad_l.cpu().numpy()

    plt.scatter(sk_f[:, 0], sk_f[:, 1], 5, sk_l, marker='.')
    plt.scatter(cad_f[:, 0], cad_f[:, 1], 5, cad_l, marker='x')
    plt.savefig(filename)
    plt.show()


def save_tsne_single_result(features, labels, filename):
    assert (len(features[:, 0]) == len(features[:, 1]))
    assert (len(features) == len(labels))

    with torch.no_grad():
        features = tsne(features, 2, 50, 20.0)

    if opt.cuda:
        features = features.cpu().numpy()
    labels = labels.cpu().numpy()

    plt.scatter(features[:, 0], features[:, 1], 5, labels, marker='x')
    plt.savefig(filename)
    plt.show()


if __name__ == "__main__":
    # Get retrieval result
    print("----- Get retrieval result -----")
    res = torch.load(res_load)
    print(res)
    sk_pd = pd.read_hdf(sk_load, 'sk')
    print(sk_pd)
    cad_pd = pd.read_hdf(cad_load, 'cad')
    print(cad_pd)

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    for index, row in tqdm(res.iterrows()):
        sk_id = row['sk_id']
        sk_class = row['class']
        mAP = row['mAp']
        sk_path = sk_pd[sk_pd['id'] == sk_id].iloc[0].name
        mAP_path = os.path.join(save_dir, 'by_mAP', str(mAP))
        cls_path = os.path.join(save_dir, 'by_class', sk_class)
        if not os.path.exists(mAP_path):
            os.makedirs(mAP_path)
        if not os.path.exists(cls_path):
            os.makedirs(cls_path)
        sk_sample = '_'.join([str(index), sk_class, sk_id, sk_path[-4:]])
        shutil.copy(sk_path, os.path.join(mAP_path, sk_sample))
        shutil.copy(sk_path, os.path.join(cls_path, sk_sample))
        for cad_id in row['cad_id']:
            cad_sample = cad_pd[cad_pd['id'] == cad_id].iloc[0]
            cad_path = cad_sample.name
            cad_sample = '_'.join([str(index), cad_sample['cat'], cad_sample['id'], cad_path[-4:]])
            shutil.copy(cad_path, os.path.join(mAP_path, cad_sample))
            shutil.copy(cad_path, os.path.join(cls_path, cad_sample))

    # Get clustering figure
    print("----- Get clustering figure -----")
    pth = torch.load(os.path.join('.', 'features', 'metric_features.pth'))
    sk_train_f = pth['sk_train_f']
    sk_train_l = pth['sk_train_l']
    sk_test_f = pth['sk_test_f']
    sk_test_l = pth['sk_test_l']
    cad_train_f = pth['cad_train_f']
    cad_train_l = pth['cad_train_l']
    cad_test_f = pth['cad_test_f']
    cad_test_l = pth['cad_test_l']

    root = os.path.join('.', 'visualize', 'shrec14')
    if not os.path.exists(root):
        os.makedirs(root)

    save_tsne_joint_result(sk_train_f, sk_train_l, cad_train_f, cad_train_l, os.path.join(root, 'joint_train.png'))
    save_tsne_joint_result(sk_test_f, sk_test_l, cad_test_f, cad_test_l, os.path.join(root, 'joint_test.png'))
    save_tsne_single_result(sk_train_f, sk_train_l, os.path.join(root, 'sk_train.png'))
    save_tsne_single_result(sk_test_f, sk_test_l, os.path.join(root, 'sk_test.png'))
    save_tsne_single_result(cad_train_f, cad_train_l, os.path.join(root, 'cad_train.png'))
    save_tsne_single_result(cad_test_f, cad_test_l, os.path.join(root, 'cad_test.png'))
